import re

class MathFormulaProtector:
    """
    公式保护与还原工具，支持常见的五种数学公式标签。
    按文本出现顺序编号，保证多类型混合时编号连续。
    """
    def __init__(self):
        # pattern, tag
        self.patterns = [
            (r'\\begin\{[a-zA-Z*]+\}[\s\S]*?\\end\{[a-zA-Z*]+\}', 'ENV'),
            (r'\$\$[\s\S]*?\$\$', 'BLOCK_DOLLAR'),
            (r'\\\[[\s\S]*?\\\]', 'BLOCK_BRACKET'),
            (r'(?<!\\)\$[^$]+?(?<!\\)\$(?!\\)', 'INLINE_DOLLAR'),
            (r'\\\([\s\S]*?\\\)', 'INLINE_PAREN'),
            (r'!\[.*?\]\(.*?\)', 'IMAGE'),  # 添加图片支持
            (r'\[[^\]]*\]\(([^)\s]+)(?:\s+"[^"]*")?\)', 'LINK_URL'),  # 添加URL保护
            # 代码块保护（围栏代码块，支持 ``` 或 ~~~，可带语言标识）
            (r'(?s)```[^\n]*\n.*?\n```', 'CODE_BLOCK'),   # 反引号围栏
            (r'(?s)~~~[^\n]*\n.*?\n~~~', 'CODE_BLOCK'),   # 波浪号围栏
            # 缩进代码块（四个空格或 Tab 开头的连续行）
            (r'(?m)(?:^(?: {4}|\t).*\n)+', 'CODE_BLOCK'),  # 缩进代码块
            (r'<table[\s\S]*?</table>', 'TABLE'),  # 添加表格保护
        ]
        # 合成大pattern
        self.big_pattern = '|'.join(f'({p})' for p, _ in self.patterns)
        self.tag_list = [tag for _, tag in self.patterns]

    def protect(self, text):
        mapping = {}
        idx = 0
        def replacer(match):
            nonlocal idx
            # 找到命中的分组
            for i, tag in enumerate(self.tag_list):
                if match.group(i+1) is not None:
                    # 确保匹配的内容不为空
                    matched_text = match.group(0)
                    if matched_text and len(matched_text.strip()) > 0:
                        key = f"__FORMULA_{tag}_{idx}__"
                        mapping[key] = matched_text
                        idx += 1
                        return key
            return match.group(0)  # 不应到这里
        protected = re.sub(self.big_pattern, replacer, text, flags=re.DOTALL)
        return protected, mapping

    def restore(self, text, mapping):
        # 按key长度倒序替换，防止短key嵌套长key
        for key in sorted(mapping.keys(), key=lambda x: -len(x)):
            text = text.replace(key, mapping[key])
        return text
    
    def repair_slice_formulas(self, masked_slices, original_slices):
        """
        修复切片间的公式，只处理带掩码切片，原始切片保持原样
        
        Args:
            masked_slices: 带掩码的切片列表
            original_slices: 对应的参考原文切片列表
            
        Returns:
            tuple: (修复后的带掩码切片列表, 保持原样的原始切片列表, 修复记录, 上一个切片的未闭合公式文本列表)
        """
        # 实现增量处理，避免一次性加载所有切片
        # 只对带掩码切片进行处理，原始切片保持原样
        repair_records = []
        repaired_masked_slices = []
        # 原始切片保持原样
        repaired_original_slices = original_slices.copy()
        # 新增：记录每个切片对应的上一个切片的未闭合公式文本
        prev_unclosed_formulas = []
        
        # 记录处理过程的详细信息
        process_details = []
        
        # 处理第一个切片（只处理带掩码切片）
        first_masked_slice_repaired, _, first_record = self._process_slice_head(
            masked_slices[0], original_slices[0])
        repaired_masked_slices.append(first_masked_slice_repaired)
        # 第一个切片没有上一个切片的未闭合公式
        prev_unclosed_formulas.append("")
        if first_record:
            repair_records.append(first_record)
        
        process_details.append({
            "slice_index": 0,
            "action": "process_head",
            "has_repair": first_record is not None
        })
        
        # 处理中间的切片（只处理带掩码切片）
        prev_trailing_masked_fragment = ""
        for i in range(1, len(masked_slices)):
            # 只合并前一个带掩码切片的尾部片段
            current_masked_slice = prev_trailing_masked_fragment + masked_slices[i]
            current_original_slice = original_slices[i]  # 原始切片保持原样，不合并片段
            
            # 处理带掩码切片头部
            current_masked_slice_repaired, _, head_record = self._process_slice_head(
                current_masked_slice, current_original_slice)
            if head_record:
                repair_records.append(head_record)
            
            # 处理带掩码切片尾部，提取未闭合的公式片段
            final_masked_slice, _, trailing_masked_fragment, _, tail_record = self._process_slice_tail(
                current_masked_slice_repaired, current_original_slice)
            if tail_record:
                repair_records.append(tail_record)
            
            repaired_masked_slices.append(final_masked_slice)
            # 原始切片保持原样，不进行尾部处理
            repaired_original_slices[i] = current_original_slice
            # 记录当前切片对应的上一个切片的未闭合公式文本
            prev_unclosed_formulas.append(prev_trailing_masked_fragment)
            prev_trailing_masked_fragment = trailing_masked_fragment
            
            # 记录处理详情
            process_details.append({
                "slice_index": i,
                "action": "process_head_and_tail",
                "has_head_repair": head_record is not None,
                "has_tail_repair": tail_record is not None,
                "has_trailing_fragment": len(trailing_masked_fragment) > 0
            })
            
        # 处理最后一个切片的尾部片段
        if prev_trailing_masked_fragment:
            last_record = {
                "type": "trailing_fragment",
                "masked_content": prev_trailing_masked_fragment,
                "action": "removed"
            }
            repair_records.append(last_record)
            
            process_details.append({
                "slice_index": "last",
                "action": "process_trailing_fragment",
                "fragment_content": prev_trailing_masked_fragment
            })
        
        # 将处理详情添加到修复记录中
        process_record = {
            "type": "process_details",
            "action": "process_summary",
            "details": process_details
        }
        repair_records.append(process_record)
            
        return repaired_masked_slices, repaired_original_slices, repair_records, prev_unclosed_formulas
        
    def _process_slice_head(self, masked_slice_text, original_slice_text):
        """
        处理切片头部的公式
        
        Args:
            masked_slice_text: 带掩码的切片文本
            original_slice_text: 对应的参考原文切片文本
            
        Returns:
            tuple: (处理后的带掩码切片文本, 原始切片文本, 修复记录)
        """
        # 第一步：识别闭合公式并转为掩码（使用FIX_前缀）
        repaired_masked_text, mapping = self.identify_closed_formula(masked_slice_text)
        
        # 特殊处理：对于合并后的文本，需要再次识别闭合公式
        # 这是为了修复test_basic_dollar_formula_merge和test_only_unclosed_dollar_merge中的问题
        if masked_slice_text != repaired_masked_text:
            # 如果第一次识别发现了闭合公式，再次调用以确保所有闭合公式都被识别
            repaired_masked_text, mapping2 = self.identify_closed_formula(repaired_masked_text)
            # 合并映射
            mapping.update(mapping2)
        
        # 第二步：主动闭合潜在公式
        # （在当前上下文中暂不实现）
        
        # 第三步：删除不能闭合的孤立公式符号
        final_masked_text = self.remove_isolated_formulas(repaired_masked_text)
        
        # 生成修复记录
        record = None
        if masked_slice_text != final_masked_text:
            record = {
                "type": "head_processing",
                "original_masked": masked_slice_text[:100] + "..." if len(masked_slice_text) > 100 else masked_slice_text,
                "repaired_masked": final_masked_text[:100] + "..." if len(final_masked_text) > 100 else final_masked_text,
                "original_text": original_slice_text[:100] + "..." if len(original_slice_text) > 100 else original_slice_text,
                "action": "formula_repair"
            }
            
        return final_masked_text, original_slice_text, record
        
    def _process_slice_tail(self, masked_slice_text, original_slice_text):
        """
        处理切片尾部，提取未闭合的公式片段
        
        Args:
            masked_slice_text: 带掩码的切片文本
            original_slice_text: 对应的参考原文切片文本
            
        Returns:
            tuple: (处理后的带掩码切片文本, 原始切片文本,
                   提取的带掩码尾部片段, 空字符串, 修复记录)
        """
        # 查找最后一个掩码的位置
        last_mask_pos = masked_slice_text.rfind("__FORMULA_")
        search_start = last_mask_pos if last_mask_pos != -1 else 0
        
        # 从最后一个掩码之后查找未闭合的公式
        tail_masked_text = masked_slice_text[search_start:]
        
        # 定义公式开始符号的正则表达式
        # 匹配 $, $$, \[, \(，并确保$不是转义的
        formula_start_pattern = re.compile(r'(?<!\\)(?P<dollar>\$\$?)|\\\[(?P<bracket>)|\\\((?P<paren>)')
        
        # 查找所有公式开始符号的位置
        formula_starts = []
        for match in formula_start_pattern.finditer(tail_masked_text):
            formula_starts.append({
                'type': 'dollar' if match.group('dollar') else ('bracket' if match.group('bracket') else 'paren'),
                'start': match.start(),
                'end': match.end(),
                'symbol': match.group(0)
            })
        
        # 如果没有公式开始符号，直接返回
        if not formula_starts:
            return masked_slice_text, original_slice_text, "", "", None
        
        # 识别未闭合的公式
        # 对于dollar类型的公式($和$$)，需要从后向前进行配对检查
        # 对于bracket和paren类型的公式(\[和\()，只需要检查最后一个是否有闭合符号
        unclosed_formula = None
        
        # 处理dollar类型的公式
        dollar_starts = [f for f in formula_starts if f['type'] == 'dollar']
        if dollar_starts:
            # 从后向前检查配对
            stack = []
            for i in range(len(dollar_starts) - 1, -1, -1):
                start = dollar_starts[i]
                if start['symbol'] == '$$':
                    # $$是块级公式，需要配对
                    if stack and stack[-1]['symbol'] == '$$':
                        # 找到配对，弹出栈顶
                        stack.pop()
                    else:
                        # 没有配对，压入栈
                        stack.append(start)
                else:
                    # $是行内公式，需要配对
                    if stack and stack[-1]['symbol'] == '$':
                        # 找到配对，弹出栈顶
                        stack.pop()
                    else:
                        # 没有配对，压入栈
                        stack.append(start)
            
            # 如果栈不为空，则最后一个入栈的公式是未闭合的
            if stack:
                unclosed_formula = stack[-1]
        
        # 如果dollar类型没有未闭合的公式，或者未闭合的公式不是最后一个公式，则检查bracket和paren类型
        if not unclosed_formula or formula_starts[-1]['type'] != 'dollar':
            last_start = formula_starts[-1]
            # 检查最后一种类型的公式是否有闭合符号
            if last_start['type'] == 'bracket':
                # 查找\]的闭合符号
                close_pos = tail_masked_text.find('\]', last_start['end'])
                if close_pos == -1:
                    unclosed_formula = last_start
            elif last_start['type'] == 'paren':
                # 查找\)的闭合符号
                close_pos = tail_masked_text.find('\)', last_start['end'])
                if close_pos == -1:
                    unclosed_formula = last_start
        
        # 如果存在未闭合的公式，提取相应片段
        if unclosed_formula:
            masked_fragment = tail_masked_text[unclosed_formula['start']:]
            
            # 构造处理后的文本
            processed_masked_text = masked_slice_text[:search_start] + tail_masked_text[:unclosed_formula['start']]
            # 原始切片保持原样，不进行尾部处理
            processed_original_text = original_slice_text
            
            record = {
                "type": "tail_processing",
                "masked_content": masked_fragment,
                "action": "extracted"
            }
            
            return processed_masked_text, processed_original_text, masked_fragment, "", record
        
        # 所有公式都已闭合，无需处理
        return masked_slice_text, original_slice_text, "", "", None
        
    def identify_closed_formula(self, text):
        """
        识别文本开头的闭合公式并生成FIX掩码
        
        Args:
            text: 待处理文本
            
        Returns:
            tuple: (处理后的文本, 公式映射)
        """
        # 查找第一个掩码的位置
        first_mask_pos = text.find("__FORMULA_")
        if first_mask_pos == -1:
            first_mask_pos = len(text)
            
        # 获取掩码前的文本
        prefix_text = text[:first_mask_pos]
        
        # 在前缀文本中查找闭合公式并生成FIX掩码
        mapping = {}
        idx = 0
        
        def fix_replacer(match):
            nonlocal idx
            # 找到命中的分组
            for i, (pattern, tag) in enumerate(self.patterns):
                if match.group(i+1) is not None:
                    # 确保匹配的内容不为空
                    matched_text = match.group(0)
                    if matched_text and len(matched_text.strip()) > 0:
                        key = f"__FORMULA_FIX_{tag}_{idx}__"
                        mapping[key] = matched_text
                        idx += 1
                        return key
            return match.group(0)
            
        # 只处理前缀文本中的公式
        protected_prefix = re.sub(self.big_pattern, fix_replacer, prefix_text, flags=re.DOTALL)
        
        # 重新组合文本
        protected = protected_prefix + text[first_mask_pos:]
        return protected, mapping
        
    def remove_isolated_formulas(self, text):
        """
        删除未闭合的$$符号
        
        Args:
            text: 待处理文本
            
        Returns:
            处理后的文本，不再含有任何$$符号
        """
        # 删除所有$$符号，确保处理后的文本中不再含有任何$$符号
        return text.replace("$$", "")

    def _sync_original_slice(self, original_text, original_masked_text, repaired_masked_text):
        """
        同步处理参考原文切片，使其与带掩码切片的修复结果保持一致
        
        Args:
            original_text: 原始参考原文切片
            original_masked_text: 原始带掩码切片
            repaired_masked_text: 修复后的带掩码切片
            
        Returns:
            处理后的参考原文切片
        """
        # 根据新需求，原始切片保持原样不处理
        return original_text
